'''
使用civit下载的模型绘图
'''
from service.painter_interface import PainterInterface
from entity.config import Config
import logging
import torch
from diffusers import StableDiffusionXLPipeline

class CivitService(PainterInterface):
    def __init__(self, model_name, model_path, hy_params, paint_params):
        self.model_name = model_name
        self.model_path = model_path
        self.hy_params = hy_params
        self.paint_params = paint_params
        logging.basicConfig(level=logging.INFO)
        self.logger = logging.getLogger(__name__)

    def paint(self, prompt: str, negative_prompt: str = None):
        self.logger.info(f"类：{self.__class__.__name__} 模型地址：{self.model_path} 模型名称：{self.model_name} 开始绘图")
        try:
            m_hy_param = {
                "torch_dtype": torch.float16,
                "safety_checker": None,
                "feature_extractor": None,
                "requires_safety_checker": False,
                "add_watermarker": False,
                "local_files_only": True
            }

            if self.hy_params:
                m_hy_param.update(self.hy_params)

            pipe = StableDiffusionXLPipeline.from_single_file(
                self.model_path, **m_hy_param
            )

            # 优化设置
            pipe = pipe.to("cuda")
            pipe.enable_attention_slicing()

            try:
                pipe.enable_xformers_memory_efficient_attention()
                self.logger.info("已开启xformers")
            except Exception as e:
                self.logger.error("开启xformers失败")

            pipe.enable_model_cpu_offload()

            images = []
            # 生成配置个数的种子
            seeds = [int(i) for i in torch.randint(0, 1000000000, (Config.get_seed_count(),)).tolist()]

            m_paint_params = {
                "num_inference_steps":50,
                "guidance_scale":7.5,
                "width": 512,
                "height": 768,
            }
            if self.paint_params:
                m_paint_params.update(self.paint_params)

            for seed in seeds:
                with torch.no_grad():
                    image = pipe(
                        prompt=prompt,
                        negative_prompt=negative_prompt,
                        generator=torch.Generator("cuda").manual_seed(seed),
                        **m_paint_params
                    ).images[0]
                    images.append(image)

            # 返回PIL Image对象
            return images
        except Exception as e:
            self.logger.error(f"类：{self.__class__.__name__} 模型地址：{self.model_path} 模型名称：{self.model_name} 绘图失败")
            self.logger.error(e)
            return None